#########################################################################
#			cluster_manager.py : 										#
#				Management class that coordinates for the cluster		#
#########################################################################

#!/usr/bin/env python

# Python imports
import os
import time
import re
import random
import xmlrpclib
import xml.marshal.generic

# eChk imports
import pm
import vm
import ss
import networkutils
import log
import siglmutils

# Cluster manager class
class ClusterManager:
	
	# Constructor
	def __init__(self):
	
		# "Defines"
		self.pingFreq = 2
		self.pingCount = 3
		self.snapshotFreq = 300  
		self.enumerationFreq = 300

		# Global vars
		self.pmList = []
		
		# XML-RPC vars
		self.node_server_port = '9000'
		
		# PM file vars
		self.pmFileName = "pmlist.xml"
	
	
	# Reads PM list from file
	def ReadPMsFromFile(self, fileName):
		 
		# Open the file
		pmFile = open(fileName, 'r')
		
		# Read the list of PM addresses
		pmAddrList = xml.marshal.generic.load(pmFile)
		
		# Close the file
		pmFile.close()
		
		# Convert the PM address list into PM objects
		pmList = []
		for pmAddr in pmAddrList:
			
			# Create new pm object
			myPM = pm.PM()
			
			# Set
			myPM.IP = pmAddr
			
			# Append
			self.pmList.append(myPM)

	# Writes PM list to a file
	def WritePMsToFile(self, fileName):
	
		# Convert the PM object list into addresses
		pmAddrList = []
		for myPM in self.pmList:
			
			# Append
			pmAddrList.append(myPM.IP)
	
		# Open the file
		pmFile = open(fileName, 'w')
		
		# Write the list
		xml.marshal.generic.dump(pmAddrList, pmFile)
		
		# Close the file
		pmFile.close()
	
	# Checks if a PM's VMs are alive
	def CheckPMAvailability(self, myPM):

		# Check if the PM is up
		if not networkutils.ping(myPM.IP):
			return True, []
		
		# Check if the VMs are all up
		downVMs = []
		for thisVM in myPM.VMs:
		
			# Ping up to pingCount times
			myPingCount = self.pingCount
			while not networkutils.ping(myPM.IP) and --myPingCount > 0:
				trash = 0
			
			# If we didn't get a response, it's down
			if myPingCount == 0:
				
				# Increment the down count
				thisVM.downCount += 1
				
				# Add to downVMs list
				downVMs.add(thisVM)
		
		# Return
		return False, downVMs
		
	# Migrates a list of VMs
	def MigrateVMList(self, downVMList):
		
		# Make sure we have snapshots for all VMs
		migrateList = []
		for downVM in downVMList:
			
			# Print a notice if we don't have a snapshot for this VM
			if downVM.GetLastSnapshot() == None:
				log.write( 'VM ' + downVM.IP + ' is down but has no recent snapshots. Unable to migrate!')
			
			# Otherwise mark it for migration
			else:
				migrateList.append(downVM)
		
		# Make sure we have a list of VMs left
		if len(migrateList) == 0:
			return
		
		# Query SigLM
		sigLMRecommendations = siglmutils.GetHosts(migrateList, self.pmList)

		"""
		# PrintSigLM results
		log.write('SigLM output:')
		for key, val in sigLMRecommendations.items():
			log.write(key + ' -> ' + val)
		"""
		
		# Migrate all VMs
		for downVM in migrateList:
			
			# Look up the PM target name
			targetPMName = sigLMRecommendations[downVM.domainCfgVars['name']]
			
			# Get the targetPM from the PM list
			targetPM = None
			for thisPM in self.pmList:
				if thisPM.IP == targetPMName:
					targetPM = thisPM
					break
			
			# Save the source
			sourcePM = downVM.myPM
			
			# Migrate VM and save target pm
			log.write(  'Target selected : ' + targetPM.IP )
			log.write(  'Migrating.' )
			self.MigrateVM(downVM, targetPM)
						
			# Reenumerate the target PM
			self.EnumeratePM(targetPM)
			
			# Reenumerate the source PM, if it isn't dead - make sure this check works
			if sourcePM in self.pmList:
				self.EnumeratePM(sourcePM)

	# Migrates a single VM
	def MigrateVM(self, myVM, targetPM):
		
		# Set state
		myVM.state = "MIGRATING"
		
		# Pickle the vm
		pickledVM = vm.PickleVM(myVM)
		
		# Connect to the target PM's XML RPC server
		xmlRPCServer = xmlrpclib.ServerProxy('http://' + targetPM.IP + ':' + self.node_server_port)
		
		# Migrate the VM to the target PM
		migrateSuccess = xmlRPCServer.Migrate(pickledVM)
		
		# Print success / failure
		if migrateSuccess == True:
			log.write( 'Migrate successfully initiated.')
		else:
			log.write( 'Migrate unsuccessful.')
		
	# Enumerates all VMs on a PM
	def EnumeratePM(self, myPM):
	
		# Print
		log.write( 'Enumerating VMs for PM: ' + myPM.IP)
		
		# Connect to the PM's XML-RPC server
		xmlRPCServer = xmlrpclib.ServerProxy('http://' + myPM.IP + ':' + self.node_server_port)
		
		# Enumerate the VMs and get marshalled result
		pickledVMList = xmlRPCServer.EnumerateVMs()
		
		# Clear the old VM list
		myPM.VMs = []
		
		# Unpickle and save the results
		for pickledVM in pickledVMList:
			
			# Unpickle
			myVM = vm.UnPickleVM(pickledVM)
			
			# Set the PM pointer
			myVM.myPM = myPM
			
			# Save on PM's VM list
			myPM.VMs.append(myVM)
		
		# Set the lastEnumerated timestamp on the PM
		myPM.lastEnumerated = time.localtime()
		
		# Print all VMs
		log.write('------------------------')
		log.write( 'PM ' + myPM.IP + ' has ' + str(len(myPM.VMs)) + ' VMs.')
		for thisVM in myPM.VMs:
			thisVM.Print()
		log.write('------------------------')
	
	# Enumerates all VMs on a PM
	def SetupSiglmPM(self, myPM):
	
		# Print
		log.write( 'Setting up SigLM for : ' + myPM.IP)
		
		# Connect to the PM's XML-RPC server
		xmlRPCServer = xmlrpclib.ServerProxy('http://' + myPM.IP + ':' + self.node_server_port)
		
		# Setup SigLM on the PM 
		xmlRPCServer.SetupSiglm()
	
	# Performs a snapshot on a VM
	def TakeVMSnapshot(self, myVM):
		
		# Debug
		log.write( "Taking snapshot for VM " + myVM.IP )
		
		# Set state
		myVM.state = "SNAPSHOTTING"
		
		# Connect to the PM's XML-RPC server
		xmlRPCServer = xmlrpclib.ServerProxy('http://' + myVM.myPM.IP + ':' + self.node_server_port)
	
		# Pickle this vm
		pickledVM = vm.PickleVM(myVM)
	
		# Perform the snapshot
		pickledSnapshot = xmlRPCServer.TakeSnapshot(pickledVM)

		# Unpickle the snapshot
		mySnapshot = ss.UnPickleSnapshot(pickledSnapshot)

		# Make sure we got something
		if mySnapshot != None:
			
			# Update this VM's snapshot list
			myVM.snapshotList.append(mySnapshot)
			
			# Set state
			myVM.state = "RUNNING"
			
			# Log
			log.write( "Completed snapshot for VM " + myVM.IP + ' at ' + mySnapshot.GetTimestampString() )
		
		# Otherwise log error
		else:
			log.write( "Unable to complete snapshot for VM " + myVM.IP )
	
	# Initialization function
	def Initialize(self):
	
		# Read PM list from file
		self.ReadPMsFromFile(self.pmFileName);
		
		# Print
		log.write( "Read " + str(len(self.pmList)) + " PMs from file.")
		
		# Query all PMs and enumerate their VMs
		for thisPM in self.pmList: 
		
			# Setup SigLM on PM
			self.SetupSiglmPM(thisPM)
			
			# Enumerate this PM's VMs
			self.EnumeratePM(thisPM)
		
	# Main function
	def Run(self):
	
		# Print
		log.write( 'Cluster manager starting.')
		log.write( 'Heartbeat frequency: ' + str(self.pingFreq) + ' seconds.')
		log.write( 'Enumeration frequency: ' + str(self.enumerationFreq) + ' seconds.')
		log.write( 'Snapshot frequency: ' + str(self.snapshotFreq) + ' seconds.')

		# Initialize
		self.Initialize()
		
		# Main program loop
		while(1):
		
			# Sleep
			time.sleep(self.pingFreq)
			
			# Create lists to hold downed PMs/VMs
			downVMList = []
			
			# Iterate over all PMs and check for downed PMs/VMs
			for thisPM in self.pmList:
			
				# Check availability
				downPM, thisPMdownVMs = self.CheckPMAvailability(thisPM)
				
				# If we have a downed PM, notify
				if downPM:
				
					# Print
					log.write( 'PM ' + thisPM.IP + ' went down. Migrating all VMs.')
					
					# Remove from the pm list
					self.pmList.remove(thisPM)
				
					# Save all its VMs in the down'd VMs list
					downVMList.extend(thisPM.VMs)
				
				# If we have some downed VMs (who must have crashed since the PM isn't down)
				elif len(thisPMdownVMs) != 0:
					
					# Print them
					for downVM in thisPMdownVMs:
						log.write( 'VM ' + downVM.IP + ' went down.')
				
					# Save them in the list
					downVMList.extend(thisPMdownVMs)
				
				# Otherwise print debug
				else:
					log.write( 'PM ' + thisPM.IP + ' and its VMs are all alive.')
				
			# If we have some downed VMs
			if len(downVMList) > 0:
				
				# Migrate them all
				self.MigrateVMList(downVMList)
			
			# Re-enumerate all PMs that have become stale
			for thisPM in self.pmList:
				
				# Check and see if we should enumerate
				if thisPM.NeedsEnumerating(self.enumerationFreq):
					
					# Enumerate
					self.EnumeratePM(thisPM)
			
			# Take snapshots for all PMs
			for thisPM in self.pmList:
			
				# Iterate over all VMs
				for thisVM in thisPM.VMs:
					
					# Check and see if we should snapshot
					if thisVM.NeedsSnapshotting(self.snapshotFreq):
						
						# Snapshot
						self.TakeVMSnapshot(thisVM)

